from numpy import (
    ndarray,  # ndarray 类型
    array,  # 创建数组
    zeros,  # 创建全零数组
    arange,  # 创建范围数组
    pi,  # 圆周率
    exp,  # 计算指数
    cos,  # 计算余弦
    arcsin,  # 计算反正弦
)

from src.utils.spherical_to_cartesian_coordinates import (
    spherical_to_cartesian_coordinates,  # 球面坐标转笛卡尔坐标的函数
)
from src.utils.euclidian_distance import (
    euclidian_distance,  # 计算欧几里得距离的函数
)
from src.utils.get_wavelength import (
    get_wavelength,  # 获取波长的函数
)


class Satellite:
    """
    Satellite 类表示一个卫星对象，包含卫星的属性和方法。
    """

    def __init__(
            self,
            rng,
            idx,
            spherical_coordinates: ndarray,
            antenna_nr: int,
            antenna_distance: float,
            antenna_gain_linear: float,
            freq: float,
            center_aod_earth_deg: float,
    ) -> None:
        """
        初始化 Satellite 类的实例。

        参数:
        - rng: 随机数生成器
        - idx: 卫星索引
        - spherical_coordinates: 球面坐标 (半径、倾斜角、离开角)
        - antenna_nr: 天线数量
        - antenna_distance: 天线间距（米）
        - antenna_gain_linear: 天线增益（线性值）
        - freq: 频率（赫兹）
        - center_aod_earth_deg: 地球中心AOD角度（度）
        """

        self.rng = rng  # 随机数生成器

        self.idx: int = idx  # 卫星索引
        self.spherical_coordinates: ndarray = spherical_coordinates  # 球面坐标
        self.cartesian_coordinates: ndarray = spherical_to_cartesian_coordinates(spherical_coordinates)  # 笛卡尔坐标

        self.antenna_nr: int = antenna_nr  # 天线数量
        self.antenna_distance: float = antenna_distance  # 天线间距（米）
        self.antenna_gain_linear: float = antenna_gain_linear  # 天线增益（线性值）

        self.freq: float = freq  # 频率（赫兹）
        self.wavelength: float = get_wavelength(self.freq)  # 波长

        self.center_aod_earth_deg: float = center_aod_earth_deg  # 地球中心AOD角度（度）

        self.distance_to_users = None  # 用户距离的字典 {用户索引: 距离}
        self.aods_to_users = None  # 用户的AOD角度（弧度）
        self.steering_vectors_to_users = None  # 用户的引导矢量（每个天线一个值）

        self.channel_state_to_users: ndarray = array([])  # 信道状态矩阵
        self.erroneous_channel_state_to_users: ndarray = array([])  # 带误差的信道状态矩阵

    def update_position(
            self,
            spherical_coordinates,
    ) -> None:
        """
        更新卫星的位置。

        参数:
        - spherical_coordinates: 新的球面坐标
        """

        self.spherical_coordinates = spherical_coordinates  # 更新球面坐标
        self.cartesian_coordinates = spherical_to_cartesian_coordinates(spherical_coordinates)  # 更新笛卡尔坐标

    def calculate_distance_to_users(
            self,
            users: list,
    ) -> None:
        """
        计算每个用户与卫星之间的距离。

        参数:
        - users: 用户对象列表
        """
        # TODO: 处理可能消失的用户

        if self.distance_to_users is None:
            self.distance_to_users = zeros(len(users))  # 初始化用户距离

        for user in users:
            self.distance_to_users[user.idx] = euclidian_distance(self.cartesian_coordinates,
                                                                  user.cartesian_coordinates)  # 计算欧几里得距离

    def calculate_aods_to_users(
            self,
            users: list
    ) -> None:
        """
        计算每个用户的离开角（AOD）。

        参数:
        - users: 用户对象列表
        """
        # AOD计算公式：
        # AOD = arcsin(
        #     ((orbit + radius_earth)^2 + sat_user_dist^2 - radius_earth^2)
        #     /
        #     (2 * (orbit + radius_earth) * sat_user_dist)
        # )
        # TODO: 处理可能消失的用户

        if self.aods_to_users is None:
            self.aods_to_users = zeros(len(users))  # 初始化AOD角度

        user_pos_idx = arange(0, len(users)) - (len(users) - 1) / 2  # 用户位置索引

        for user in users:
            # 计算每个用户的AOD
            self.aods_to_users[user.idx] = arcsin(
                (
                        + self.spherical_coordinates[0] ** 2
                        + self.distance_to_users[user.idx] ** 2
                        - user.spherical_coordinates[0] ** 2
                )  # 分子
                /
                (
                        2 * self.spherical_coordinates[0] * self.distance_to_users[user.idx]
                )  # 分母
            )

            # 如果用户在中心AOD的右侧，调整AOD值
            if user_pos_idx[user.idx] >= 0:
                self.aods_to_users[user.idx] = 2 * (self.center_aod_earth_deg * pi / 180) - self.aods_to_users[user.idx]

    def calculate_steering_vectors(
            self,
            users: list,
    ) -> None:
        """
        计算每个用户的引导矢量（对于给定的ULA和AOD）。

        参数:
        - users: 用户对象列表
        """

        if self.steering_vectors_to_users is None:
            self.steering_vectors_to_users = zeros((len(users), self.antenna_nr), dtype='complex64')  # 初始化引导矢量

        steering_idx = arange(0, self.antenna_nr) - (self.antenna_nr - 1) / 2  # 天线索引

        for user in users:
            # 计算引导矢量
            self.steering_vectors_to_users[user.idx, :] = exp(
                steering_idx * (
                    -1j * 2 * pi / self.wavelength
                    * self.antenna_distance
                    * cos(self.aods_to_users[user.idx])
                )
            )

    def update_channel_state_information(
            self,
            channel_model,
            users: list,
    ) -> None:
        """
        更新每个用户的信道状态信息。

        参数:
        - channel_model: 信道模型函数
        - users: 用户对象列表
        """
        # 使用给定的信道模型更新信道状态信息
        self.channel_state_to_users = channel_model(self, users)

    def update_erroneous_channel_state_information(
            self,
            error_model_config,
            users: list,
    ) -> None:
        """
        更新每个用户的带有误差的信道状态信息。

        参数:
        - error_model_config: 误差模型配置
        - users: 用户对象列表
        """
        # 使用给定的误差模型更新带有误差的信道状态信息
        self.erroneous_channel_state_to_users = error_model_config.error_model(error_model_config=error_model_config,
                                                                               satellite=self,
                                                                               users=users)
